#include "llvm/Transforms/MPI/MPI.h"
#include "llvm/Transforms/Scalar.h"
#include "llvm/Intrinsics.h"
#include "llvm/LLVMContext.h"
#include "llvm/Module.h"

#include "llvm/Support/IRBuilder.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/Target/TargetData.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Config/config.h"

using namespace llvm;

namespace {
  class MPIAnnotateFunctions : public ModulePass {
  private:
    void setDoesNotThrow(Function &F);
    void setDoesNotCapture(Function &F, unsigned n);
    void setDoesNotAlias(Function &F, unsigned n);

    bool Modified;
  public:
    static char ID; // Pass identification, replacement for typeid
    MPIAnnotateFunctions() : ModulePass(ID) {}

    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
      ModulePass::getAnalysisUsage(AU);
      AU.setPreservesAll();
    }

    virtual bool runOnModule(Module &M);
  };
}

char MPIAnnotateFunctions::ID = 0;
static RegisterPass<MPIAnnotateFunctions>
X("mpi-annotate", "Annotate MPI Functions with AA/ModRef information.");
//INITIALIZE_PASS(MPIAnnotateFunctions, "mpi-annotate", "Annotate MPI Functions with AA/ModRef information.", false, false);

// Publically exposed interface to pass...
const PassInfo *const MPIAnnotateFunctionsID = &X;
ModulePass *createMPIAnnotateFunctionsPass() { return new MPIAnnotateFunctions(); }

// Copied and adapted from SimplifyLibCalls
void MPIAnnotateFunctions::setDoesNotThrow(Function &F) {
  if (!F.doesNotThrow()) {
    F.setDoesNotThrow();
    Modified = true;
  }
}
void MPIAnnotateFunctions::setDoesNotCapture(Function &F, unsigned n) {
  if (!F.doesNotCapture(n)) {
    F.setDoesNotCapture(n);
    Modified = true;
  }
}
void MPIAnnotateFunctions::setDoesNotAlias(Function &F, unsigned n) {
  if (!F.doesNotAlias(n)) {
    F.setDoesNotAlias(n);
    Modified = true;
  }
}

bool MPIAnnotateFunctions::runOnModule(Module &M) {
  StringRef MPIPrefix("MPI_");

  Modified = false;
  for (Module::iterator I = M.begin(), E = M.end(); I != E; ++I) {
    Function &F = *I;
    if (!F.isDeclaration())
      continue;

    if (!F.hasName())
      continue;

    const FunctionType *FTy = F.getFunctionType();

    StringRef Name = F.getName();
    if(!Name.startswith(MPIPrefix))
      continue;


    // Since this doesn't define MOD/REF behavior, we can treat
    // ISend and IRecv the same way.
    if (Name == "MPI_Isend" || Name == "MPI_Irecv") {
      // They both have the same signature, so we can verify as either one.
      if (!MPISupport::verifyIrecvSignature(FTy, &F.getContext()))
        continue;
      setDoesNotThrow(F);

      // TODO: Can we say that they do not capture buf?
      if(isa<PointerType>(FTy->getParamType(MPIArgs::IDatatype)))
        setDoesNotCapture(F, MPIArgs::IDatatype);

      if(isa<PointerType>(FTy->getParamType(MPIArgs::IComm)))
        setDoesNotCapture(F, MPIArgs::IComm);

      // Request is newly allocated on return
      setDoesNotAlias(F, MPIArgs::IRequest);

    } else if (Name == "MPI_Wait") {
      if (!MPISupport::verifyWaitSignature(FTy, &F.getContext()))
        continue;

      setDoesNotThrow(F);
      setDoesNotAlias(F, MPIArgs::WaitRequest);
      setDoesNotCapture(F, MPIArgs::WaitStatus);
    }
  }

  return Modified;
}
